import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
from sklearn import model_selection
from sklearn import metrics

dataset = pd.read_csv('../data/watermelon-3a.csv')
#数据预处理
X = dataset[['density','sugar']]
Y = dataset['classification']
#分割训练集和验证集
X_train,X_test,Y_train,Y_test = model_selection.train_test_split(X,Y,test_size=0.5,random_state=0)
print(X_train, Y_train)
#训练
LDA_model = LinearDiscriminantAnalysis()
LDA_model.fit(X_train,Y_train)
#验证
Y_pred = LDA_model.predict(X_test)
#汇总
print(metrics.confusion_matrix(Y_test, Y_pred))
print(metrics.classification_report(Y_test, Y_pred, target_names=['Bad','Good']))
print(LDA_model.coef_)

w = [LDA_model.coef_[0,0], LDA_model.coef_[0,1]]
print(w)
#画图
good_melon = dataset[dataset['classification'] == 1]
bad_melon = dataset[dataset['classification'] == 0]
plt.scatter(bad_melon['density'],bad_melon['sugar'],marker='o',color='r',s=100,label='bad')
plt.scatter(good_melon['density'],good_melon['sugar'],marker='o',color='g',s=100,label='good')

x_line = np.arange(0, 1, 0.05)
y_line = -w[0] / w[1] * x_line
plt.plot(x_line, y_line)
plt.show()